package com.wangjie.prefixsum;

import java.util.HashMap;

/**
 * @author jieshao
 * @date 2022/3/25 11:41
 *
 * 560. 和为 K 的子数组
 */
public class Test03 {
    public static void main(String[] args) {
        int[] arr = {1, 2, 3, 0, 6};
        int k = 6;
        System.out.println("------------方式1------------");
        int result1 = subarraySum1(arr, k);
        System.out.println("result1 = " + result1);

        System.out.println("------------方式2------------");
        int result2 = subarraySum2(arr, k);
        System.out.println("result2 = " + result2);
    }

    public static int subarraySum1(int[] num, int k) {
        int n = num.length;
        // 构造前缀和
        int[] preSum = new int[n + 1];
        for (int i = 0; i < n; i++) {
            preSum[i + 1] = preSum[i] + num[i];
        }
        int res = 0;
        // 穷举所有子数组
        for (int i = 0; i < n; i++) {
            for (int j = i + 1; j <= n; j++) {
                if (preSum[j] - preSum[i] == k) {
                    res++;
                }
            }
        }
        return res;
    }

    public static int subarraySum2(int[] num, int k) {
        int n = num.length;
        // preSum 前缀和 -> 该前缀和出现的次数
        HashMap<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        // 和为 k 的子数组个数
        int res = 0;
        // 前缀和
        int sum0_i = 0;
        for (int i = 0; i < n; i++) {
            sum0_i += num[i];
            // 这是我们想找的前缀和 num[0..j]
            int sum0_j = sum0_i - k;
            // 如果前面有这个前缀和，则直接更新答案
            if (map.containsKey(sum0_j)) {
                res += map.get(sum0_j);
            }
            // 把前缀和 num[0..i] 加入并记录出现次数
            map.put(sum0_i, map.getOrDefault(sum0_i, 0) + 1);
        }

        return res;
    }
}